# coding: UTF-8
import time
import torch
from torch.utils.data import DataLoader
import numpy as np
from train_eval import train, init_network
from importlib import import_module
import argparse
from dataset.CMN_dataset import CMNDataset
from utils import get_time_dif
parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', default='DPCNN', type=str, required=False, help='choose a model: TextCNN, TextRNN, FastText, TextRCNN, TextRNN_Att, DPCNN, Transformer')
#parser.add_argument('--embedding', default='random', type=str, help='random or pre_trained')
#parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')
args = parser.parse_args()

if __name__ == '__main__':
    dataset = 'source/data/multi_classification'  # 数据集
    model_name = args.model  # 'TextRCNN'  # TextCNN, TextRNN, FastText, TextRCNN, TextRNN_Att, DPCNN, Transformer
    

    x = import_module('models.' + model_name)
    config = x.Config(dataset)
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True

    start_time = time.time()
    print("Loading data...")
    train_dataset = CMNDataset(config, 'train', rebuild_vocab=False)
    val_dataset = CMNDataset(config, 'val')
    
    # train
    train_dataloader = DataLoader(
        train_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    val_dataloader = DataLoader(
        val_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)
    model1 = x.Model(config).to(config.device)
    model2 = x.Model(config).to(config.device)
    if model_name != 'Transformer':
        init_network(model1)
        init_network(model2)
    print(model1.parameters)
    train(config, model1, model2, train_dataloader, val_dataloader)
